#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import math
from pathlib import Path
from typing import List, Dict

import torch
from torch import Tensor
from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer

def render(pc, vp_cam, bg_color, scaling_modifier=1.0) -> Dict[str, Tensor]:

    # Set up rasterization configuration
    raster_settings = GaussianRasterizationSettings(
                    image_height=vp_cam.image_height,
                    image_width=vp_cam.image_width,
                    tanfovx=math.tan(vp_cam.FoVx * 0.5),
                    tanfovy=math.tan(vp_cam.FoVy * 0.5),
                    bg=bg_color,
                    scale_modifier=scaling_modifier,
                    viewmatrix=vp_cam.world_view_transform,
                    projmatrix=vp_cam.full_proj_transform,
                    sh_degree=0,
                    campos=vp_cam.camera_center,
                    prefiltered=False,
                    debug=pc.hp.debug,
                )
    rasterizer = GaussianRasterizer(raster_settings=raster_settings)

    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
    screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device='cuda')
    
    try: screenspace_points.retain_grad()
    except: pass

    means3D = pc.get_xyz
    means2D = screenspace_points
    opacity = pc.get_opacity
    scales = pc.get_scaling
    rotations = pc.get_rotation
    
    visible_mask = rasterizer.markVisible(positions=means3D)   # (N, 1)
    colors_precomp = pc.feature_encoder(vp_cam, visible_mask)
    mask = (opacity[visible_mask] >= 1/255).squeeze(-1)
    # Rasterize visible Gaussians to image, obtain their radii (on screen).
    rendered_image, radii = rasterizer(
        means3D=means3D[visible_mask][mask],
        means2D=means2D[visible_mask][mask],
        shs=None,
        colors_precomp=colors_precomp[mask],
        opacities=opacity[visible_mask][mask],
        scales=scales[visible_mask][mask],
        rotations=rotations[visible_mask][mask],
        cov3D_precomp=None)
    
    view_radii = torch.zeros_like(mask, dtype=radii.dtype, device="cuda")
    view_radii[mask] = radii
    tem_radii = torch.zeros_like(visible_mask, dtype=radii.dtype, device="cuda")
    tem_radii[visible_mask] = view_radii
    visibility_filter = torch.logical_and(visible_mask, tem_radii>0)
              
    return {
        'render': rendered_image,
        'viewspace_points': screenspace_points,
        'visibility_filter': visibility_filter,
        'radii': tem_radii,
        'opacity_visible': visible_mask}
